from ppgnss import gnss_utils
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.axes_grid1 import make_axes_locatable

import utils

plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']

width, height = gnss_utils.cm2inch(18), gnss_utils.cm2inch(8)

fontsize=10
filename = "span5_comp_va_5min_down.txt"
data = utils.read_data_file(filename)
nrow = data.shape[0]
data = data.head(nrow-2)
color = "gray"
color = "magma_r" # "RdYlGn_r"
x = np.arange(len(data["Utide"]))
y1 = [0, 1]
yy1, xx1 = np.meshgrid(y1, x)
snr_thred = 2
# 创建图形和坐标轴
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(width, height), gridspec_kw={'height_ratios': [1, 1.2]})
z = np.zeros((len(data["Utide"]), len(y1)))
for i, _ in enumerate(["S", "D"]):
    z[:, i] = data[f"{_}amp"]
    mask = data[f"{_}snr"] < snr_thred
    z[:, i][mask] = np.nan
# z = z.reshape(1, len(z))
# axes[0].pcolormesh(np.arange(1, 3), np.arange(1, 11), matrix, cmap='viridis', edgecolors='k', linewidth=0.5)
im = axes[0].pcolormesh(xx1, yy1, z, edgecolors="k", linewidth=0.2, cmap=plt.colormaps['magma_r'])
axes[0].set_xticks(x)
axes[0].set_xticklabels(data["Utide"], fontsize=fontsize, rotation=270)
axes[0].set_yticks([0, 1])
axes[0].set_yticklabels(["闸外水位", "内外位差"], fontsize=fontsize)
axes[0].set_xlabel(u'潮汐分量', fontsize=fontsize)
divider = make_axes_locatable(axes[0])
cax = divider.new_vertical(size = '15%', pad = 0.3)
fig.add_axes(cax)
cbar = fig.colorbar(im, cax = cax, orientation = 'horizontal')
ztklbls = np.arange(0, 2.6, 0.5)
cbar.set_ticks(ztklbls)  # 根据刻度标签的数量设置刻度位置
cbar.set_ticklabels(["%3.1f m" % _ for _ in ztklbls])

ax = axes[1]
y = [0, 1, 2]
yy, xx = np.meshgrid(y, x)
z = np.zeros((len(data["Utide"]), len(y)))
# z[:, 0] = data["Uamp"]
# z[:, 1] = data["Yamp"]
# z[:, 2] = data["Xamp"] 
for icomp, comp in enumerate(["U", "Y", "X"]):
    z[:, icomp] = data[f"{comp}amp"].values*100
    mask = data[f"{comp}snr"] < snr_thred
    z[:, icomp][mask] = np.nan

plt.subplots_adjust(left=0.1, right=0.95, bottom=0.22, top=0.98, hspace=1)
colors = cm.Greens(z / z.max())  # 使用viridis颜色映射，可以根据需要选择其他颜色映射
# 绘制3D柱状图
# 设置柱状图的尺寸
dx,dy = 0.001, 0.1  # 在 x、y 轴方向上的宽度
dz = z  # 在 z 轴方向上的高度
print(xx.shape, yy.shape, z.shape)
im=ax.pcolormesh(xx, yy, z,  edgecolors="k", linewidth=0.2, cmap=plt.colormaps['RdYlGn_r'])
ax.set_xticks(x)
ax.set_xticklabels(data["Utide"], fontsize=fontsize, rotation=270)
ax.set_yticks(y)
ax.set_yticklabels(["U", "Y", "X"], fontsize=fontsize)
# 添加标签和标题
ax.set_xlabel(u'潮汐分量', fontsize=fontsize)
ax.set_ylabel(u'坐标分量', fontsize=fontsize)
# ax.set_title('3D Bar Chart with Color Fill')
# cbar = plt.colorbar(cax, orientation='horizontal')
# add color bar below chart
divider = make_axes_locatable(ax)
cax = divider.new_vertical(size = '10%', pad = 0.3)
fig.add_axes(cax)
cbar = fig.colorbar(im, cax = cax, orientation = 'horizontal')
ztklbls = np.arange(0, 1.5, 0.2)
cbar.set_ticks(ztklbls)  # 根据刻度标签的数量设置刻度位置
cbar.set_ticklabels(["%3.1f mm" % _ for _ in ztklbls])
    # 显示图形
fig_filename = "figures/colored_freqs.png"
plt.savefig(fig_filename, dpi=300)
plt.show()